import requests
import json
import time


class AiMedal():
    def __init__(self):
        self.API_KEY = "khTseCrI3yB0X4YuisRLgd1w"
        self.SECRET_KEY = "lozK9Yu6aFwN7FFOooUhHxoecx73fZyg"
    
    def submit_task(self, prompt):
        """
        提交AI作画任务
        :param prompt: 用户输入的提示词
        :return: 包含成功状态和任务ID的字典
        """
        print(f"提交AI作画任务，提示词: {prompt}")
        try:
            url = "https://aip.baidubce.com/rpc/2.0/wenxin/v1/basic/textToImage?access_token=" + self.get_access_token()

            payload = json.dumps({
                "text": prompt,
                "resolution": "512*512"
            }, ensure_ascii=False)
            headers = {
                'Accept': 'application/json',
                'Content-Type': 'application/json'
            }

            print(f"请求URL: {url}")
            print(f"请求payload: {payload}")

            response = requests.post(url, headers=headers, data=payload.encode("utf-8"))
            print(f"响应状态码: {response.status_code}")
            print(f"响应内容: {response.text}")

            response_data = response.json()
            
            if response.status_code == 200 and 'data' in response_data and 'taskId' in response_data['data']:
                task_id = response_data['data']['taskId']
                print(f"任务提交成功，任务ID: {task_id}")
                return {
                    'success': True,
                    'task_id': task_id,
                    'message': '任务提交成功'
                }
            else:
                error_msg = response_data.get('error_msg', '任务提交失败')
                error_code = response_data.get('error_code', '')
                print(f"任务提交失败: {error_msg} (错误代码: {error_code})")
                
                # 根据错误代码提供更友好的错误信息
                if error_code == 6:
                    friendly_msg = "API权限不足，请检查百度AI服务是否已开通文心一言图像生成功能"
                elif error_code == 282501:
                    friendly_msg = "提示词长度超出限制，请简化描述"
                else:
                    friendly_msg = f"AI服务暂时不可用: {error_msg}"
                
                return {
                    'success': False,
                    'message': friendly_msg,
                    'error_code': error_code,
                    'original_error': error_msg
                }
                
        except Exception as e:
            print(f"提交任务异常: {str(e)}")
            return {
                'success': False,
                'message': f'提交任务失败: {str(e)}'
            }

    def get_access_token(self):
        """
        使用 AK，SK 生成鉴权签名（Access Token）
        :return: access_token，或是None(如果错误)
        """
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {"grant_type": "client_credentials", "client_id": self.API_KEY, "client_secret": self.SECRET_KEY}
        return str(requests.post(url, params=params).json().get("access_token"))

    def main(self):
        url = "https://aip.baidubce.com/rpc/2.0/wenxin/v1/basic/textToImage?access_token=" + self.get_access_token()

        payload = json.dumps({
            "text": "生成一只暴躁白色兔兔的勋章",
            "resolution": "512*512"
        }, ensure_ascii=False)
        headers = {
            'Accept': 'application/json'
        }

        response = requests.request("POST", url, headers=headers, data=payload.encode("utf-8"))

        print(response.text)


class GetMedal():
    def __init__(self):
        self.API_KEY = "khTseCrI3yB0X4YuisRLgd1w"
        self.SECRET_KEY = "lozK9Yu6aFwN7FFOooUhHxoecx73fZyg"
    
    def get_result(self, task_id):
        """
        获取AI作画结果
        :param task_id: 任务ID
        :return: 包含成功状态和图片数据的字典
        """
        print(f"获取AI作画结果，任务ID: {task_id}")
        try:
            url = "https://aip.baidubce.com/rpc/2.0/wenxin/v1/basic/getImg?access_token=" + self.get_access_token()

            # 尝试不同的task_id格式
            payload = json.dumps({
                "taskId": str(task_id)  # 确保task_id是字符串格式
            }, ensure_ascii=False)
            headers = {
                'Accept': 'application/json',
                'Content-Type': 'application/json'
            }

            print(f"请求URL: {url}")
            print(f"请求payload: {payload}")

            response = requests.post(url, headers=headers, data=payload.encode("utf-8"))
            print(f"响应状态码: {response.status_code}")
            print(f"响应内容: {response.text}")

            response_data = response.json()
            
            if response.status_code == 200 and 'data' in response_data:
                if 'status' in response_data['data']:
                    status = response_data['data']['status']
                    print(f"任务状态: {status}")
                    
                    if status == 1:  # 任务完成
                        # 检查是否有图片URL
                        if 'img' in response_data['data'] and response_data['data']['img']:
                            image_data = response_data['data']['img']
                            print(f"成功获取图片URL: {image_data}")
                            return {
                                'success': True,
                                'image_data': image_data,
                                'message': '生成成功'
                            }
                        elif 'imgUrls' in response_data['data'] and response_data['data']['imgUrls']:
                            # 检查第一个图片是否被阻止
                            first_img = response_data['data']['imgUrls'][0]
                            if first_img.get('img_approve_conclusion') == 'block':
                                print("图片被内容审核阻止")
                                return {
                                    'success': False,
                                    'message': '图片被内容审核阻止，请尝试使用更安全的描述'
                                }
                            elif first_img.get('image'):
                                image_data = first_img['image']
                                print(f"成功获取图片URL: {image_data}")
                                return {
                                    'success': True,
                                    'image_data': image_data,
                                    'message': '生成成功'
                                }
                            else:
                                print("图片URL为空")
                                return {
                                    'success': False,
                                    'message': '图片生成失败，请重试'
                                }
                        else:
                            print("响应数据中没有图片URL")
                            return {
                                'success': False,
                                'message': '图片生成失败，请重试'
                            }
                    elif status == 0:  # 任务进行中
                        print(f"任务进行中，状态: {status}")
                        return {
                            'success': False,
                            'message': f'任务进行中，状态: {status}',
                            'retry': True  # 添加retry标志，表示需要重试
                        }
                    else:
                        print(f"未知任务状态: {status}")
                        return {
                            'success': False,
                            'message': f'未知任务状态: {status}'
                        }
                else:
                    # 直接返回图片数据的情况
                    if 'images' in response_data['data']:
                        image_data = response_data['data']['images'][0]['url']
                        print(f"直接获取图片URL: {image_data}")
                        return {
                            'success': True,
                            'image_data': image_data,
                            'message': '生成成功'
                        }
                    else:
                        print("响应数据格式异常")
                        return {
                            'success': False,
                            'message': '响应数据格式异常'
                        }
            else:
                error_msg = response_data.get('error_msg', '获取结果失败')
                error_code = response_data.get('error_code', '')
                print(f"获取结果失败: {error_msg}, 错误代码: {error_code}")
                
                # 处理特定的错误情况
                if error_code == '282004':  # task_id无效
                    return {
                        'success': False,
                        'message': '任务ID无效，可能是任务还在处理中，请稍后重试',
                        'retry': True
                    }
                elif error_code == '501':  # 生图失败
                    return {
                        'success': False,
                        'message': '图片生成失败，可能是内容审核未通过，请尝试使用更安全的描述',
                        'retry': False
                    }
                elif 'task_id' in error_msg.lower() or '无效' in error_msg:
                    return {
                        'success': False,
                        'message': '任务ID无效，请稍后重试',
                        'retry': True
                    }
                elif '生图失败' in error_msg:
                    return {
                        'success': False,
                        'message': '图片生成失败，请尝试使用更安全的描述或稍后重试',
                        'retry': False
                    }
                else:
                    return {
                        'success': False,
                        'message': error_msg
                    }
                
        except Exception as e:
            print(f"获取结果异常: {str(e)}")
            return {
                'success': False,
                'message': f'获取结果失败: {str(e)}'
            }

    def get_access_token(self):
        """
        使用 AK，SK 生成鉴权签名（Access Token）
        :return: access_token，或是None(如果错误)
        """
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {"grant_type": "client_credentials", "client_id": self.API_KEY, "client_secret": self.SECRET_KEY}
        return str(requests.post(url, params=params).json().get("access_token"))

    def main(self):
        url = "https://aip.baidubce.com/rpc/2.0/wenxin/v1/basic/getImg?access_token=" + self.get_access_token()

        payload = json.dumps({
            "taskId": "1958526669180620761"
        }, ensure_ascii=False)
        headers = {
            'Accept': 'application/json'
        }

        response = requests.request("POST", url, headers=headers, data=payload.encode("utf-8"))

        print(response.text)


# 创建实例
aiMedal = AiMedal()
getMedal = GetMedal()




